package com.aikg.kgrag_java.pojo.entity;

import com.aikg.kgrag_java.pojo.dto.neo4j.Neo4jRetrievalDTO;

import java.util.List;
import java.util.Map;

/**
 * @ClassName: Prompts
 * @Description: prompt类，get各种prompt
 * @Author: 李宗雨
 * @Date: 2025/1/3 16:57
 * @Version: 1.0
 **/
public class Prompts {

    private final static String RAG_PROMPT = "以下是`结构化数据`:\n" +
                                            "`实体及其属性列表`:\n" +
                                            "{node_info_list}\n" +
                                            "`关系列表`:\n" +
                                            "```\n"+
                                            "{relation_info_list}\n" +
                                            "```\n\n" +
                                            "以下是`相关文本片段`:\n" +
                                            "{context}\n\n" +
                                            "请严格按照以上`结构化数据`和`相关文本片段`回答以下问题，**多余信息不要输出**。\n" +
                                            "问题:{query}";

    private final static String QUESTION_RECOMMEND_PROMPT_1 =
            "以下是{n}个段落：\n" +
            "{related_docs}" +
            "问题:{question}\n" +
            "答案:{answer}\n\n" +
            "请你判断问题答案出自哪个段落。仅返回段落序号，其他多余信息不要返回。";

    private final static String QUESTION_RECOMMEND_PROMPT_2 =
            "上下文信息：{doc}\n"+
            "根据提供的上下文信息和没有先验知识的原则，仅基于以下查询生成问题。\n\n"+
            "你是一名教师/教授。你的任务是围绕上下文信息设置{k}个主题相关问题。这些问题的答案必须在上下文信息中找得到。仅输出三对问题和答案，其他多余信息不要输出。";

    public static String getRagPrompt(Neo4jRetrievalDTO neo4jRetrievalDTO,
                                      String context,
                                      String query
    ){
        return RAG_PROMPT.replace("{node_info_list}", neo4jRetrievalDTO.nodeInfoListToString4Me())
                .replace("{relation_info_list}", neo4jRetrievalDTO.relationTripletListToString4Me())
                .replace("{context}", context)
                .replace("{query}", query);

    }

    public static String getQuestionRecommendPrompt1(ConversationInfo conversationInfo) {
        StringBuilder tmpStr = new StringBuilder();
        for(int i=1; i <= conversationInfo.getChunks().size(); i++){
            tmpStr.append("段落")
                    .append(i)
                    .append(":")
                    .append(conversationInfo.getChunks().get(i - 1).getPage_content())
                    .append("\n\n");
        }
        return QUESTION_RECOMMEND_PROMPT_1
                .replace("{n}",conversationInfo.getChunks().size() + "")
                .replace("{related_docs}", tmpStr.toString())
                .replace("{query}", conversationInfo.getQuery())
                .replace("{answer}", conversationInfo.getAnswer());
    }

    /**
     * 制作生成问题的prompt
     *
     * @param chunk 最相关文本块
     * @param k 生成k个问题
     * @return 生成问题的prompt
     */
    public static String getQuestionRecommendPrompt2(String chunk, int k) {
        return QUESTION_RECOMMEND_PROMPT_1
                .replace("{doc}", chunk)
                .replace("{k}", k+"");
    }

}
